#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# # author : cypro666
# # date   : 2015.08.01
# # wrapper of Decision Tree method in sklearn
import sys, json, time
import numpy as np
from threading import Thread
from sklearn import tree

from magic3.utils import Timer
from magic3.filesystem import *
add_sys_path(grand_dir(__file__))

from skt.base import MethodBase
from skt.utils import _normalize, make_options


class DecisionTree(MethodBase):
    def __init__(self, parameters={}):
        super().__init__(parameters)

    @property
    def name(self): return 'DecisionTree'

    def output(self, cls):
        self.Log('output')
        super().output()
        extra = {}

        try:    extra['num_classes'] = [int(i) for i in cls.n_classes_]
        except: extra['num_classes'] = int(cls.n_classes_)

        try:    extra['max_features'] = int(cls.max_features_)
        except: pass

        try:    extra['classes'] = [int(i) for i in cls.classes_]
        except: pass

        try:    extra['feature_importance'] = [float(i) for i in cls.feature_importances_]
        except: pass

        fn = self.json_name()
        json.dump(extra, open(fn, 'w'), indent=4)

    def execute(self):
        self.read_input()
        self._normalize(False, self._param['normalize'])

        if self._param['min_samples_split'] not in set(range(1, 100)):
            self._param['min_samples_split'] = 2

        if self._param['min_samples_leaf'] not in set(range(1, 10)):
            self._param['min_samples_leaf'] = 1

        if self._param['min_weight_fraction_leaf'] >= 1.0:
            self._param['min_weight_fraction_leaf'] = 0.0

        self.save_parameters()

        cls = tree.DecisionTreeClassifier(criterion="gini",
                                          splitter="best",
                                          min_samples_split=self._param['min_samples_split'],
                                          min_samples_leaf=self._param['min_samples_leaf'],
                                          min_weight_fraction_leaf=self._param['min_weight_fraction_leaf'],
                                          max_features=None,
                                          max_depth=None,
                                          max_leaf_nodes=None,
                                          random_state=None,
                                          class_weight=None)

        cls.fit(self._train, self._label)

        self._results = cls.predict(self._testing)
        self.output(cls)

    def run(self, timeout):
        t = Thread(target=self.execute)
        t.start()
        t.join(timeout)
        if t.is_alive():
            self.Log('timeout!')
        self.Log('exit')

if __name__ == '__main__':
    opts = [('train_file', 'str', []),
            ('label_file', 'str', []),
            ('testing_file', 'str', []),
            ('results_file', 'str', []),
            ('log_file', 'str', []),
            ('normalize', 'choice', ['0', '1', '2']),
            ('min_samples_split', 'int', []),
            ('min_samples_leaf', 'int', []),
            ('min_weight_fraction_leaf', 'float', []),
            ('max_features', 'choice', ['auto', 'sqrt', 'log2'])]

    parameters = make_options(opts)

    timer = Timer()

    DecisionTree(parameters).run(300)

    if __debug__:
        print('elapsed:', timer)


